import tensorflow.compat.v1 as tf
import numpy as np
import matplotlib.pyplot as plt
import uuid

tf.disable_v2_behavior()

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def computeYByX(x):
    noise = np.random.normal(-50, 50, x.shape)
    return 400 * np.sin(x) + 2 * x * x + noise


xTrain = np.linspace(-20, 20, 401).reshape([1, -1])
noise = np.random.normal(-0.2, 0.2, xTrain.shape)
yTrain = computeYByX(xTrain)

# 保存初始训练数据图片
plt.clf()
plt.plot(xTrain[0], yTrain[0], 'ro', label=u'训练数据')
plt.legend()
plt.savefig('curve_data.png', dpi=200)

x = tf.placeholder(tf.float32, shape=[1, 401])


def weight_variable(shape):
    initial = tf.truncated_normal(shape, stddev=0.1)
    return tf.Variable(initial)


def bias_variable(shape):
    initial = tf.constant(0.1, shape=shape)
    return tf.Variable(initial)


hiddenDim = 400

W = weight_variable([hiddenDim, 1])
b = bias_variable([hiddenDim, 1])

W2 = weight_variable([1, hiddenDim])
b2 = bias_variable([1])

W3 = weight_variable([401, 401])
b4 = bias_variable([1])

hidden = tf.nn.sigmoid(tf.matmul(W, x) + b)
y = tf.matmul(W2, hidden) + b2

loss = tf.reduce_mean(tf.square(y - yTrain))
step = tf.Variable(0, trainable=False)
rate = tf.train.exponential_decay(0.15, step, 1, 0.9999)
optimizer = tf.train.AdamOptimizer(rate)
train = optimizer.minimize(loss, global_step=step)
init = tf.global_variables_initializer()

sess = tf.Session()
sess.run(init)

for time in range(0, 10001):
    train.run({x: xTrain}, sess)
    if time % 1000 == 0:
        print('训练次数：', time, ',训练损失平均值：', loss.eval({x: xTrain}, sess))

        plt.clf()
        plt.plot(xTrain[0], yTrain[0], 'ro', label=u'训练数据')
        plt.plot(xTrain[0], y.eval({x: xTrain}, sess)[0], label=u'拟合曲线')
        plt.legend()
        plt.savefig('curve_fitting_' + str(int(time / 1000)) + '.png', dpi=200)

xTest = np.linspace(-40, 40, 401).reshape([1, -1])
yTest = computeYByX(xTest)

plt.clf()
plt.plot(xTrain[0], yTrain[0], 'mo', label=u'训练数据')
plt.plot(xTrain[0], y.eval({x: xTrain}, sess)[0], label=u'拟合曲线')
plt.legend()
plt.savefig(str(uuid.uuid4()) + '.png', dpi=200)
plt.show()
